package tianhao.luo.huffmancode;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.*;

/**
 * @author: tianhao.luo@hand-china.com 2021/6/26  上午8:32
 */
public class HuffmanCode {

    Map<Byte, String> codes = new HashMap<>();
    /**
     * 创建带数据域和权值的Node
     * @param bytes 数据域
     * @return 结点list
     */
    private List<Node> getNodes(byte[] bytes){
        List<Node> nodes = new ArrayList<>();
        Map<Byte, Integer> map = new HashMap<>();
        // 存储每一个Byte出现的次数
        for (byte aByte : bytes) {
            Integer count = map.get(aByte);
            if (count == null) {
                map.put(aByte,1);
            } else {
                count++;
                map.put(aByte,count);
            }
        }
        // 遍历map,转化为Node
        for (Map.Entry<Byte, Integer> entry : map.entrySet()) {
            nodes.add(new Node(entry.getKey(),entry.getValue()));
        }
        return nodes;
    }

    /**
     * 创建哈夫曼树
     * @param list node结点list
     * @return 哈夫曼树的根节点
     */
    private Node createHuffmanTree(List<Node> list){

        while (list.size()>1){
            // 从小到大排序
            Collections.sort(list);
            // 取出第一颗最小的二叉树
            Node leftNode = list.get(0);
            Node rightNode = list.get(1);
            // 创建新的二叉树,它没有data,只有权值
            Node parent = new Node(null, leftNode.weight + rightNode.weight);
            parent.left = leftNode;
            parent.right = rightNode;

            // 将处理过的二叉树从nodes删除
            list.remove(leftNode);
            list.remove(rightNode);

            // 将新的二叉树,加入list
            list.add(parent);
        }
        // list中最后的结点,就是哈弗曼树的根节点
        return list.get(0);
    }

    public void preOrder(Node node){
        if (node !=null){
            node.preOrder();
        }else {
            System.out.println("哈弗曼树为空");
        }
    }

    // 生成赫夫曼树对应的哈弗曼编码
    // 1.将哈弗曼编码表放在Map<Byte,String>中
    // 2.生成哈弗曼编码时,需要拼接路径,就是叶子节点的路径

    /**
     * 将传入的node结点的所有叶子节点的哈弗曼编码的到,并放入huffmanCodes集合中
     * @param node 传入结点
     * @param code 路径:左子节点0,右子节点1
     * @param stringBuilder 用于拼接路径
     * @param map 存放叶子节点对应的哈弗曼编码
     */
    private void getCodes(Node node,String code,StringBuilder stringBuilder,Map<Byte,String > map){
        StringBuilder stringBuilder1 = new StringBuilder(stringBuilder);
        // 将code加入到StringBuilder1,根节点存放""
        stringBuilder1.append(code);
        // 结点不为空才可以遍历
        if (node!=null){
            if (node.data == null){
                // 非叶子节点,需要递归处理
                // 向左
                getCodes(node.left,"0",stringBuilder1,map);
                // 向右
                getCodes(node.right,"1",stringBuilder1,map);
            }else {
                // 叶子节点,表示找到了
                map.put(node.data,stringBuilder1.toString());
            }
        }
    }

    public Map<Byte,String> getCodes(Node root){
        // 考虑根节点为空的情况
        if (root == null){
            return null;
        }
        Map<Byte, String> map = new HashMap<>(10);
        StringBuilder stringBuilder = new StringBuilder();
        getCodes(root,"",stringBuilder,map);
        return map;
    }

    // 将字符串对应的byte数组,通过生成的哈弗曼编码表,返回一个哈弗曼编码压缩后的byte[]

    /**
     *
     * @param bytes 原始字符串对应的byte数组
     * @param huffmanCodes 哈弗曼编码
     * @return 返回哈弗曼编码处理后的byte数组
     * 例如:"i like like like java do you like a java"
     *
     */
    private byte[] zip(byte[] bytes,Map<Byte,String> huffmanCodes){

        // 1.利用huffmanCodes将bytes转成哈弗曼编码对应的字符串
        StringBuilder stringBuilder = new StringBuilder();
        for (byte aByte : bytes) {
            stringBuilder.append(huffmanCodes.get(aByte));
        }

//        System.out.println(stringBuilder.toString());
        // 此时这个StringBuilder还无法用于传递,因为将原字符串转化为哈弗曼编码对应的字节的字符串描述后,反而让字符串长度增加了
        // 所以需要StringBuilder转化为byte数组

        // 统计返回的字节数组长度
        int length;
        if (stringBuilder.length() % 8 ==0){
            length = stringBuilder.length()/8;
        }else {
            length = stringBuilder.length()/8 +1;
        }

        // 创建一个存储压缩后的byte数组
        byte[] temp =new byte[length];
        int index=0;
        // 因为每8位才对应一个byte,所以步长+8
        for (int i = 0; i < stringBuilder.length(); i+=8) {
            String strByte;
            if (i+8 > stringBuilder.length()){
                // 最后剩余的不够八位
                strByte = stringBuilder.substring(i);
            }else {
                strByte = stringBuilder.substring(i, i + 8);
            }
            // 将strByte转化为byte数组
            temp[index] = (byte) Integer.parseInt(strByte,2);
            index++;
        }
        return temp;
    }

    /**
     * 封装压缩方法
     * @param bytes 字节数组
     * @return 压缩后的字节数组
     */
    public byte[] huffmanZip(byte[] bytes){
        List<Node> nodes = getNodes(bytes);
        Node huffmanTree = createHuffmanTree(nodes);
        codes = getCodes(huffmanTree);
        return zip(bytes, codes);
    }

    /**
     * 讲一个byte转化为一个二进制字符串
     * @param flag 表示是否需要补高位,如果为true表示需要补高位,如果为false则不需要
     * @param b 字节
     * @return 二进制字符串
     */
    private String bytesToBitString(boolean flag,byte b){
        // 保存b
        int temp =b;
        // 如果是正数,还需要补高位
        // temp 1==>00000001
        // 按位与256, 1 0000 0000  | 0000 0001 => 1 0000 0001
        if (flag){
            temp |= 256;
        }
        // 返回的是temp二进制的补码
        String s = Integer.toBinaryString(temp);
        // 需要截取补码的后八位
        if (flag){
            return s.substring(s.length()-8);
        }else {
            return s;
        }
    }


    // 对压缩数据解码

    /**
     *
     * @param huffmanCodes 哈弗曼编码表
     * @param huffmanBytes 哈弗曼编码表处理后的字节数组
     * @return 原来的字符串对应的数组
     */
    private byte[] decode(Map<Byte,String> huffmanCodes,byte[] huffmanBytes){
        // 1.先得到huffmanBytes对应的二进制字符串
        StringBuilder stringBuilder = new StringBuilder();
        // 2.将byte数组转化为二进制的字符串
        for (int i = 0; i < huffmanBytes.length; i++) {
            // 判断是否为最后一个字节
            boolean flag = (i != huffmanBytes.length -1);
            stringBuilder.append(bytesToBitString(flag,huffmanBytes[i]));
        }

        // 将二进制字符串按照哈弗曼编码进行解码
        // 将哈弗曼编码表进行调换,要反向查询
        Map<String, Byte> map = new HashMap<>();
        for (Map.Entry<Byte, String> entry : huffmanCodes.entrySet()) {
            map.put(entry.getValue(),entry.getKey());
        }

        // 创建一个集合,存放byte
        List<Byte> list = new ArrayList<>();
        StringBuilder tempBuilder = new StringBuilder();
        for (int i = 0; i < stringBuilder.length(); i++) {
            tempBuilder.append(stringBuilder.substring(i,i+1));
            if (map.containsKey(tempBuilder.toString())){
                list.add(map.get(tempBuilder.toString()));
                tempBuilder.delete(0,tempBuilder.length());
            }
        }
        byte[] tempBytes = new byte[list.size()];
        for (int i = 0; i < tempBytes.length; i++) {
            tempBytes[i] = list.get(i);
        }
        return tempBytes;
    }

    public byte[] decode(byte[] huffmanBytes,byte[] bytes){
        List<Node> nodes = getNodes(bytes);
        Node huffmanTree = createHuffmanTree(nodes);
        Map<Byte, String> codes = getCodes(huffmanTree);
        return decode(codes, huffmanBytes);
    }


    // 编写方法,将一个文件进行压缩

    /**
     *
     * @param srcFile 待压缩的文件目录
     * @param dstFile 压缩好的文件目录
     */
    public void zipFile(String srcFile,String dstFile){


        // 创建输出流
        // 创建文件的输入流
        try(FileInputStream fileInputStream = new FileInputStream(srcFile);
            FileOutputStream fileOutputStream = new FileOutputStream(dstFile)
        ) {
            // 创建一个和源文件一样大小的字节数组
            byte[] b = new byte[fileInputStream.available()];
            // 读取文件
            fileInputStream.read(b);

            // 使用哈弗曼编码
            byte[] bytes = huffmanZip(b);
            // 创建和文件输出流相关的ObjectOutputStream
            final ObjectOutputStream objectOutputStream = new ObjectOutputStream(fileOutputStream);

            objectOutputStream.writeObject(bytes);
            // 注意需要将哈弗曼编码也写入压缩文件
            // 以对象流方式写入哈弗曼编码,是为了恢复源文件使用
            objectOutputStream.writeObject(codes);
        }catch (Exception e){
            System.out.println(e.getMessage());
        }
    }


    // 编写一个方法完成对压缩文件的解压

    /**
     *
     * @param zipFile 准备解压的文件
     * @param destFile 将文件解压到的路径
     */
    public void unZipFIle(String zipFile,String destFile){
        try(final FileInputStream fileInputStream = new FileInputStream(zipFile);
            final FileOutputStream fileOutputStream = new FileOutputStream(destFile)) {

            final ObjectInputStream objectInputStream = new ObjectInputStream(fileInputStream);
            byte[] bytes = (byte[]) objectInputStream.readObject();
            // 读取哈弗曼编码表
            Map<Byte, String> huffmanCodes = (Map<Byte, String>) objectInputStream.readObject();
            // 解码
            final byte[] decode = decode(huffmanCodes, bytes);
            // 写入目标文件
            fileOutputStream.write(decode);
        }catch (Exception e){
            System.out.println(e.getMessage());
        }
    }
}
